#!/bin/bash
#
# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
# Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

set -e

WDIR=
IMAGE_VERSION=
KERNEL_VERSION=
CHROOT_DIR=
INITRD=
VMLINUZ=
MODULES_DIR=

SIGNING_CERT=
SIGNING_KEY=
CDIR=/sonic/
SDIR=$CDIR/platform/${CONFIGURED_PLATFORM}/installer/
DDIR=$CDIR
CAPSULE=/lib/firmware/mellanox/boot/capsule/boot_update2.cap
BOOTCTL_DRIVER=
TMFIFO_DRIVER=
SDHCI_OF_DWCMSHC_DRIVER=
TARGET_MACHINE=nvidia-bluefield
BFB="${BFB:-/lib/firmware/mellanox/boot/default.bfb}"
GRUB_AA64=grubnetaa64.efi
GRUB_CFG="" # Common Grub Config
BF2_BOOT_ARGS="console=ttyAMA1 console=hvc0 console=ttyAMA0 earlycon=pl011,0x01000000 earlycon=pl011,0x01800000"
BF2_GRUB_CFG="$BF2_BOOT_ARGS isolcpus=1-7 nohz_full=1-7 rcu_nocbs=1-7"
BF3_BOOT_ARGS="console=ttyAMA1 console=hvc0 console=ttyAMA0 earlycon=pl011,0x13010000"
BF3_GRUB_CFG="$BF3_BOOT_ARGS isolcpus=1-13 nohz_full=1-13 rcu_nocbs=1-13"

usage() {
cat << EOF
Usage: `basename $0` [ OPTIONS ]
OPTIONS:
-k, --kernel <kernel version>     Kernel version for the SmartNIC.
-sc, --signing-cert               Secure upgrade signing certificate.
-sk, --signing-key                Secure upgrade signing key.
-v, --verbose                     Run script in verbose mode. Will print out each step of execution.
-h, --help                        Display help
EOF
}
parse_args() {
    while [[ $@ != "" ]]; do
        case $1 in
            -h|--help)
                usage
                exit 0
                ;;
            -v|--verbose)
                shift
                set -x
                ;;
            -k|--kernel)
                shift
                KERNEL_VERSION=$1
                ;;
        -sc|--signing-cert)
            shift
        echo "signing cert $1"
        SIGNING_CERT="$1"
        ;;
            -sk|--signing-key)
        shift
        echo "signing key $1"
        SIGNING_KEY="$1"
        ;;
            *)
                usage
                exit 1
                ;;
        esac
        shift
    done
}

clean_dir()
{
    rm -rf $1
    exit $2
}

validate_config() {
    if [[ ! -f "$INITRD" ]]; then
        echo "[create_sonic_image] Error! SONiC INITRD not found"
        exit 1
    fi

    if [[ ! -f "$VMLINUZ" ]]; then
        echo "[create_sonic_image] Error! SONiC VMLINUZ not found"
        exit 1
    fi

    if [[ ! -d "$MODULES_DIR" ]]; then
        echo "[create_sonic_image] Error! Path to Kernel Modules not found"
        exit 1
    fi

    if [[ ! -n "$OUTPUT_BFB_IMAGE" ]]; then
        echo "[create_sonic_image] Error! OUTPUT_BFB_IMAGE name not defined. Exiting.."
        exit 1
    fi

    if [[ ! -f "$CDIR/$INSTALLER_PAYLOAD" ]]; then
        echo "$INSTALLER_PAYLOAD not found. Exiting.."
        exit 1
    fi

    mkbfb=`which mlx-mkbfb`
    if [[ ! -x "${mkbfb}" ]]; then
        echo "Error! mlx-mkbfb is required to build BFB image"
        exit 1
    fi

    if [[ ! -f $CAPSULE ]]; then
        echo "ERROR: Capsule file $CAPSULE does not exist"
        exit 1
    fi

    if [[ ! -f $BOOTCTL_DRIVER ]]; then
        echo "ERROR: Bootctl driver $BOOTCTL_DRIVER does not exist"
        exit 1
    fi

    if [[ ! -f $TMFIFO_DRIVER ]]; then
        echo "ERROR: Bootctl driver $TMFIFO_DRIVER does not exist"
        exit 1
    fi

    if [[ ! -f $SDHCI_OF_DWCMSHC_DRIVER ]]; then
        echo "ERROR: Bootctl driver $SDHCI_OF_DWCMSHC_DRIVER does not exist"
        exit 1
    fi

    if [[ ! -f $BFB ]]; then
        echo "ERROR: Default BFB $BFB does not exist"
        exit 1
    fi
}

cleanup_workdir() {
    rm -rf $boot_args $boot_args2 $boot_path $boot_desc $WDIR
}

create_workdir() {
    if [[ $WDIR == "" ]]; then
        WDIR=$(mktemp -d /sonic/bfb-wd-XXXX)
    else
        rm -rf $WDIR/*
    fi
    boot_args=$(mktemp /tmp/boot-args-XXXX)
    boot_args2=$(mktemp /tmp/boot-args2-XXXX)
    boot_path=$(mktemp /tmp/boot-path-XXXX)
    boot_desc=$(mktemp /tmp/boot-desc-XXXX)
    trap cleanup_workdir EXIT
}

add_sonic_to_initramfs() {
    # Add the logic to put second stage installer into bfb initramfs
    mkdir -p debian

    j2 ${SDIR}/install.sh.j2 -o ./debian/install.sh -e SECURE_UPGRADE_MODE="$SECURE_UPGRADE_MODE"
    chmod 0755 ./debian/install.sh

    # Copy the INSTALLER payload
    cp $CDIR/$INSTALLER_PAYLOAD ./debian/

cat > scripts/initrd-install << EOF
#!/bin/bash

depmod -a $KERNEL_VERSION > /dev/null 2>&1
insmod /mlx-bootctl.ko
insmod /sdhci-of-dwcmshc.ko
insmod /sbsa_gwdt.ko
/usr/sbin/watchdog

/bin/bash /debian/install.sh

EOF
    chmod +x scripts/initrd-install
    sudo rm -f conf/conf.d/debian-core*
    sudo sh -c 'echo "BOOT=initrd-install" > conf/conf.d/initrd_install'
    sudo sh -c 'echo "ROOT=\"LABEL=writable\"" > conf/conf.d/default_root'
}

copy_bin()
{
    if [ -e $1 ]; then
        bin=$1
    else
        bin=$(which $1 2> /dev/null)
    fi
    if [ -z "$bin" ]; then
        echo "ERROR: Cannot find $1"
        exit 1
    fi
    sudo mkdir -p .$(dirname $bin)
    if [ ! -e .${bin} ]; then
        sudo cp -a $bin .${bin}
    fi

    # Copy dependencies
    for lib in  $(ldd $bin 2> /dev/null | grep '=>' | awk '{print $3}')
    do
        if [ -e .$lib ]; then
            continue
        fi
        sudo mkdir -p .$(dirname $lib)
        sudo cp -a $lib .$lib
        if [ -h $lib ]; then
            tlib=$(readlink -f $lib)
            if [ ! -e .$tlib ]; then
                sudo mkdir -p .$(dirname $tlib)
                sudo cp $tlib .$tlib
            fi
        fi
    done
}

create_bfb_image() {

    pushd $WDIR

    # Copy the initrd into the work directory
    initramfs=$(realpath $INITRD)
    cp $initramfs $WDIR/dump-initramfs-v0
    initramfs=$WDIR/dump-initramfs-v0

    case "$(file --brief --mime-type "$initramfs")" in
        "application/x-lzma")
            cat_initrd="lzcat" ;;
        "application/x-lz4")
            cat_initrd="lz4cat" ;;
        "application/zstd")
            cat_initrd="zstdcat" ;;
        *)
            cat_initrd="zcat" ;;
    esac

    echo "Rebuilding $initramfs"

    mkdir -p ${WDIR}/initramfs
    pushd initramfs
    $cat_initrd "$initramfs" | cpio -i

    # Remove tools coming with busybox
    for tool in `dpkg -L grub2-common` \
            `dpkg -L e2fsprogs` \
            `dpkg -L kmod | grep -v share` \
            `dpkg -L pciutils | grep -v share` \
            `dpkg -L usbutils | grep -v share` \
            `dpkg -L tar` \
            /usr/sbin/watchdog
    do
        if [ -d $tool ]; then
            continue
        fi
        /bin/rm -f .${tool}
    done

    for tool in `dpkg -L bfscripts | grep bin/` \
            `dpkg -L e2fsprogs | grep -v share` \
            `dpkg -L grub2-common` \
            `dpkg -L kmod | grep -v share` \
            `dpkg -L pciutils | grep -v share` \
            `dpkg -L usbutils | grep -v share` \
            `dpkg -L tar | grep -v share` \
            `dpkg -L grub-efi-arm64-bin` \
            `dpkg -L dmidecode | grep -v share` \
            xz efibootmgr bash getopt hexdump lspci perl \
            lsblk shutdown systemctl strings aarch64-linux-gnu-strings \
            mlxbf-bootctl id mkfs fsck watchdog  dirname curl openssl
    do
        if [ -d $tool ]; then
            continue
        fi
        copy_bin $tool
    done

    kernel_mft=$(dpkg -l | grep kernel-mft-dkms-modules | awk '/^ii/ {print $2}')
    if [[ $kernel_mft == "" ]]; then
        echo "ERROR: kernel-mft-dkms-modules package is not installed"
        exit 1
    fi

    for tool in `dpkg -L mft` \
            `dpkg -L mft-oem` \
            `dpkg -L $kernel_mft` \
            `dpkg -L xmlstarlet | grep -v share`
    do
        if [ -d $tool ]; then
            continue
        fi
        copy_bin $tool
    done

    sudo depmod -a -b ./ $KERNEL_VERSION
    mkdir -p usr/share/misc/ bin/ var/log/watchdog
    sudo cp /etc/watchdog.conf etc
    sudo cp /usr/share/misc/pci.ids usr/share/misc/
    cp $CHROOT_DIR/usr/share/misc/pci.ids usr/share/misc/
    cp $BOOTCTL_DRIVER .
    cp $TMFIFO_DRIVER .
    cp $SDHCI_OF_DWCMSHC_DRIVER .
    cp $WATCHDOG .
    mkdir -p ./secure-boot
    cp -r $CDIR/$FILESYSTEM_ROOT/boot/* ./secure-boot

    mkdir -p ./lib/firmware/mellanox/boot/
    cp /lib/firmware/mellanox/boot/default.bfb ./lib/firmware/mellanox/boot/default.bfb
    cp -a /lib/firmware/mellanox/boot/capsule ./lib/firmware/mellanox/boot/
    mkdir -p mnt dev sys proc

    add_sonic_to_initramfs

    # Make initramfs with new debian
    find . -print0 | sudo cpio --null -o --format=newc | gzip -9 > "$initramfs"

    popd

    printf "$BF2_BOOT_ARGS initrd=initramfs" > \
        "$boot_args"
    printf "$BF3_BOOT_ARGS initrd=initramfs" > \
        "$boot_args2"

    printf "VenHw(F019E406-8C9C-11E5-8797-001ACA00BFC4)/Image" > "$boot_path"
    printf "Linux from rshim" > "$boot_desc"
    vmlinuz=$WDIR/vmlinuz
    cat $VMLINUZ > "$vmlinuz"

    $mkbfb --image "$vmlinuz" \
           --initramfs "$initramfs" \
           --capsule "$CAPSULE" \
           --boot-args-v0 "$boot_args" \
           --boot-args-v2 "$boot_args2" \
           --boot-path "$boot_path" \
           --boot-desc "$boot_desc" \
            ${BFB} ${DDIR}/${OUTPUT_BFB_IMAGE}

    echo "BFB is ready: $(readlink -f  ${DDIR}/${OUTPUT_BFB_IMAGE})"

    popd
}

main() {
    echo $@
    parse_args $@
    . $CDIR/onie-image-arm64.conf
    # Export ENV Variables for j2cli
    CHROOT_DIR=$CDIR/$FILESYSTEM_ROOT
    if [[ ! -d "$CHROOT_DIR" ]]; then
        echo "[create_sonic_image] Error! Path to CHROOT not found"
        exit 1
    fi
    export GRUB_CFG=$(cat $SDIR/sonic-grub.cfg)
    export IMAGE_VERSION=$(cat $CHROOT_DIR/etc/sonic/sonic_version.yml | grep "build_version" | sed -e "s/build_version: //g;s/'//g")
    export BF2_BOOT_ARGS BF3_BOOT_ARGS BF2_GRUB_CFG BF3_GRUB_CFG INSTALLER_PAYLOAD FILESYSTEM_DOCKERFS DOCKERFS_DIR FILESYSTEM_SQUASHFS KERNEL_VERSION SECURE_UPGRADE_MODE SIGNING_KEY SIGNING_CERT

    INITRD=$CDIR/$FILESYSTEM_ROOT/boot/initrd.img-$KERNEL_VERSION
    VMLINUZ=$CDIR/$FILESYSTEM_ROOT/boot/vmlinuz-$KERNEL_VERSION
    MODULES_DIR=$CDIR/$FILESYSTEM_ROOT/lib/modules/$KERNEL_VERSION
    WATCHDOG=$MODULES_DIR/kernel/drivers/watchdog/sbsa_gwdt.ko
    BOOTCTL_DRIVER=$CHROOT_DIR/usr/lib/modules/$KERNEL_VERSION/extra/mlx-bootctl.ko
    TMFIFO_DRIVER=$CHROOT_DIR/usr/lib/modules/$KERNEL_VERSION/extra/mlxbf-tmfifo.ko
    SDHCI_OF_DWCMSHC_DRIVER=$CHROOT_DIR/usr/lib/modules/$KERNEL_VERSION/extra/sdhci-of-dwcmshc.ko

    validate_config

    create_workdir
    echo "Work directory: $WDIR"

    create_bfb_image
}

main $@
